#-*- coding: utf-8 -*-
#
# Copyright 2013 Antoine Drouin (poinix@gmail.com)
#
# This file is part of PAT.
#
#    PAT is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    PAT is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with PAT.  If not, see <http://www.gnu.org/licenses/>.
#

"""
This is a 3dof dynamic model for a planar VTOL
"""

import math
import numpy as np
import scipy.integrate
import matplotlib.pyplot as plt
import pdb

import pat.dynamic_model as dm
import pat.utils as pu
import pat.frames as fr

'''
State vector components
'''
sv_x    = 0
sv_z    = 1
sv_th   = 2
sv_xd   = 3
sv_zd   = 4
sv_thd  = 5
sv_size = 6

'''
input vector components
'''
iv_f1   = 0
iv_f2   = 1
iv_size = 2

def dyn(X, t, U, param):
    """
    State Space Representation: Xdot = f_param(t, X, U)
    """
    Xd = np.zeros(sv_size)
    Xd[sv_x:sv_th+1] = X[sv_xd:sv_thd+1]

    sth = math.sin(X[sv_th])
    cth = math.cos(X[sv_th])
    ut = (U[iv_f1] + U[iv_f2])
    ud = (U[iv_f2] - U[iv_f1]) 

    Xd[sv_xd]  = -sth/param.m*ut
    Xd[sv_zd]  =  cth/param.m*ut-9.81
    Xd[sv_thd] = param.l/param.J*ud

    return Xd    


def jacobian(Xe, Ue, param):
    """
    Partial derivatives of the dynamics:
    A = d/dX(dyn)(Xe, Ue), B = d/dU(dyn)(Xe, Ue)
    """   
    st = math.sin(Xe[sv_th]); ct = math.cos(Xe[sv_th])
    one_over_m = 1/param.m 
    A = np.array([[0, 0, 0, 1, 0, 0],
                  [0, 0, 0, 0, 1, 0],
                  [0, 0, 0, 0, 0, 1],
                  [0, 0, -one_over_m*ct*(Ue[iv_f1]+Ue[iv_f2]), 0, 0, 0],
                  [0, 0, -one_over_m*st*(Ue[iv_f1]+Ue[iv_f2]),  0, 0, 0],
                  [0, 0, 0, 0, 0, 0]])
    l_over_J = param.l/param.J
    B = np.array([[0, 0],
                  [0, 0],
                  [0, 0],
                  [-one_over_m*st, -one_over_m*st],
                  [ one_over_m*ct,  one_over_m*ct],
                  [-l_over_J, l_over_J]])
    return A, B   


def trim(param, args=None):
    """
    Return equilibrium state Xe and input Ue such as
    dyn(Xe, Ue) = 0
    """
    Xe = np.zeros(sv_size)
    if args<>None:
        try: Xe[sv_x] = args['x']
        except KeyError: pass
        try: Xe[sv_z] = args['z']
        except KeyError: pass
    Ue =  0.5*param.m*9.81*np.ones(iv_size)
    return Xe, Ue
    


class Param():
    def __init__(self):
        self.m = 0.5
        self.l = 0.2
        self.J = 0.01

    def __str__(self):
        return "m:{:.3f} kg\nl:{:.2f}m\nJ:{:.4f} kg.m2".format(self.m, self.l, self.J)


class DynamicModel(dm.DynamicModel):

    sv_x     =  sv_x
    sv_z     =  sv_z
    sv_th    =  sv_th
    sv_xd    =  sv_xd
    sv_zd    =  sv_zd
    sv_thd   =  sv_thd
    sv_size  =  sv_size

    iv_f1   = iv_f1
    iv_f2   = iv_f2
    iv_size = iv_size

    dyn = lambda self, X, t, U, P: dyn(X, t, U, self.P)
    trim = lambda self, args=None: trim(self.P, args)
    jacobian = lambda self, Xe, Ue: jacobian(Xe, Ue, self.P)

#    plot_trajectory = lambda self, time, X, U=None: plot_trajectory(time, X, U)

    def __init__(self, args=None):
        dm.DynamicModel.__init__(self)
        self.X = np.zeros(sv_size)
        self.P = Param()
        self.reset()

    def name(self):
        return "planar rotorcraft python simple"

    def reset(self, X0=None):
        if X0<>None: self.X = X0
        else: self.X = np.zeros(sv_size); self.X[sv_th] += pu.rad_of_deg(1.)
        return self.X

    def run(self, dt, U):
        foo, self.X = scipy.integrate.odeint(dyn, self.X, [0, dt], args=(U, self.P, ))
        return self.X

    def param(self):
        return str(self.P)

    def plot_state(self, time, X, U=None, figure=None, window_title=None, filename=None):
        plot_trajectory(time, X, U, figure, window_title, filename=filename)
 
        
    def state_SixDOFfEuclidianEuler(self):
        fmt = fr.SixDOFfEuclidianEuler
        X = np.zeros(fmt.size)
        X[fmt.x] = self.X[sv_x]
        X[fmt.z] = -self.X[sv_z]
        X[fmt.theta] = self.X[sv_th]
        return X

    def state_str(self):
        return """pos: {:-.2f}, {:-.2f} m
        vel: {:-.2f}, {:-.2f} m/s
        att:    {:-.2f} deg
        """.format(self.X[sv_x],  self.X[sv_z],
                   self.X[sv_xd], self.X[sv_zd],
                   pu.deg_of_rad(self.X[sv_th]))


'''
TODO: fine tune sizes and margins
'''
def plot_trajectory2(time, X, U=None, figure=None, window_title="Trajectory",
                    legend=None, filename=None):
    plots = [("$x$",             "m",     X[:,sv_x]),
             ("$z$",             "m",     X[:,sv_z]),
             ("$\\theta$",       "m",     pu.deg_of_rad(X[:, sv_th])),
             ("$\dot{x}$",       "m/s",   X[:, sv_xd]),
             ("$\dot{z}$",       "m/s",   X[:, sv_zd]),
             ("$\dot{\\theta}$", "deg/s", pu.deg_of_rad(X[:, sv_th]))]
    if U<>None:
        plots += [("$U$", "N", U)]

    return pu.plot_in_grid(time, plots, 3, figure, window_title, legend, filename,
                           margins=(0.04, 0.05, 0.98, 0.96, 0.10, 0.34))
    
def plot_trajectory(time, X, U=None, figure=None, window_title="Trajectory",
                    legend=None, filename=None):
    if U<>None: nrow = 3; figsize = (20.48, 10.24)
    else: nrow=2; figsize = (20.48, 7.24)

    margins=(0.04, 0.04, 0.99, 0.95, 0.19, 0.29)
    figure = pu.prepare_fig(figure, window_title, figsize=figsize, margins=margins)

    ax = plt.subplot(nrow, 3, 1)
    plt.plot(time, X[:, sv_x])
    pu.decorate(ax, title="$x$", ylab="m", xlab='$time\; in\; s$')
    ax = plt.subplot(nrow, 3, 2)
    plt.plot(time, X[:, sv_z])
    pu.decorate(ax, title="$z$", ylab="m", xlab='$time\; in\; s$')
    ax = plt.subplot(nrow, 3, 3)
    plt.plot(time, pu.deg_of_rad(X[:, sv_th]))
    pu.decorate(ax, title="$\\theta$", ylab="deg", xlab='$time\; in\; s$')
    ax = plt.subplot(nrow, 3, 4)
    plt.plot(time, X[:, sv_xd])
    pu.decorate(ax, title="$\dot{x}$", ylab="m/s", xlab='$time\; in\; s$')
    ax = plt.subplot(nrow, 3, 5)
    plt.plot(time, X[:, sv_zd])
    pu.decorate(ax, title="$\dot{z}$", ylab="m/s", xlab='$time\; in\; s$')
    ax = plt.subplot(nrow, 3, 6)
    plt.plot(time, pu.deg_of_rad(X[:, sv_thd]))
    pu.decorate(ax, title="$\dot{\\theta}$", ylab="deg/s", xlab='$time\; in\; s$')
    if U<>None:
        ax = plt.subplot(nrow, 3, 7)
        plt.plot(time, U)
        pu.decorate(ax, title="$U$", ylab="N", legend=['F1', 'F2'])
    pu.save_if(filename)
